#define _CRT_SECURE_NO_WARNINGS
#define STRICT
#include <windows.h>
#include "callback.h"
#include "hooked.h"
#include "handling.h"
#include <cstdlib>
#include <algorithm>
using namespace std;

bool operator<(const HookedFunction& a, const HookedFunction& b)
{
	return a.originalFunction < b.originalFunction;
}

__declspec(noreturn) void ExitViolation(const char* message)
{
	strcpy(callbackInformation->backMessage.message, message);
	callbackInformation->backMessage.finishCode = BackMessage::Violation;
	ExitProcess(0);
}

__declspec(noreturn) void ExitViolationFile(const char* message, const char* fileName)
{
	sprintf_s(callbackInformation->backMessage.message, sizeof(callbackInformation->backMessage.message), "%s: %s", message, fileName);
	callbackInformation->backMessage.finishCode = BackMessage::Violation;
	ExitProcess(0);
}

DWORD ExitCrash(unsigned code)
{
	strcpy(callbackInformation->backMessage.message, "Crash by thread observer");
	callbackInformation->backMessage.finishCode = BackMessage::Crash;
	ExitProcess(code);
}

HMODULE WINAPI HookedLoadLibraryA(LPCSTR lpLibFileName)
{
	HMODULE module = LoadLibraryA(lpLibFileName);
	if(module)
		InjectToModule((char*)module);
	return module;
}

HMODULE WINAPI HookedLoadLibraryW(LPCWSTR lpLibFileName)
{
	HMODULE module = LoadLibraryW(lpLibFileName);
	if(module)
		InjectToModule((char*)module);
	return module;
}

BOOL WINAPI HookedFreeLibrary(HMODULE hLibModule)
{
	if(hLibModule == injectModule)
		ExitViolation("FreeLibrary: invalid call");
	return FreeLibrary(hLibModule);
}

__declspec(noreturn) VOID WINAPI HookedFreeLibraryAndExitThread(HMODULE hLibModule, DWORD dwExitCode)
{
	if(hLibModule == injectModule)
		ExitViolation("FreeLibraryAndExitThread: invalid call");
	FreeLibraryAndExitThread(hLibModule, dwExitCode);
}

HANDLE WINAPI HookedCreateFileA(LPCSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile)
{
	const char* allowedFiles = callbackInformation->injectSettings.allowedFilesA;
	size_t i;
	for(i = 0; allowedFiles[i]; i += strlen(allowedFiles + i) + 1)
		if(_stricmp(lpFileName, allowedFiles + i) == 0)
			break;
	if(!allowedFiles[i])
		ExitViolationFile("CreateFileA, invalid file name", lpFileName);
	return CreateFileA(lpFileName, dwDesiredAccess, FILE_SHARE_READ | FILE_SHARE_WRITE, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);
}

HANDLE WINAPI HookedCreateFileW(LPCWSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile)
{
	const wchar_t* allowedFiles = callbackInformation->injectSettings.allowedFilesW;
	size_t i;
	for(i = 0; allowedFiles[i]; i += wcslen(allowedFiles + i) + 1)
		if(_wcsicmp(lpFileName, allowedFiles + i) == 0)
			break;
	if(!allowedFiles[i])
	{
		char buf[256];
		wcstombs(buf, lpFileName, sizeof(buf) - 1);
		ExitViolationFile("CreateFileW, invalid file name", buf);
	}
	return CreateFileW(lpFileName, dwDesiredAccess, FILE_SHARE_READ | FILE_SHARE_WRITE, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);
}

static DWORD CALLBACK ThreadRoutine(LPVOID lpParameter)
{
	__try
	{
		pair<LPTHREAD_START_ROUTINE, LPVOID>* parameter = (pair<LPTHREAD_START_ROUTINE, LPVOID>*)lpParameter;
		LPTHREAD_START_ROUTINE routine = parameter->first;
		lpParameter = parameter->second;
		delete parameter;
		return routine(lpParameter);
	}
	__except(ExitCrash(GetExceptionCode()))
	{
		return 0;
	}
}

HANDLE WINAPI HookedCreateThread(LPSECURITY_ATTRIBUTES lpThreadAttributes, SIZE_T dwStackSize, LPTHREAD_START_ROUTINE lpStartAddress, LPVOID lpParameter, DWORD dwCreationFlags, LPDWORD lpThreadId)
{
	return CreateThread(lpThreadAttributes, dwStackSize, ThreadRoutine, new pair<LPTHREAD_START_ROUTINE, LPVOID>(lpStartAddress, lpParameter), dwCreationFlags, lpThreadId);
}

template <typename T>
HookedFunction MakeHookedFunction(T originalFunction, T hookFunction, unsigned hookFlags = 0)
{
	HookedFunction function;
	function.function = hookFunction;
	function.originalFunction = originalFunction;
	function.hookFlags = hookFlags;
	return function;
}

#define MHF(lib, func) MakeHookedFunction(func, Hooked##func)
#define MHFF(lib, func, flags) MakeHookedFunction(func, Hooked##func, flags)

static HookedFunction functions[] =
{
	MHF(kernel32, LoadLibraryA),
	MHF(kernel32, LoadLibraryW),
	MHF(kernel32, FreeLibrary),
	MHF(kernel32, FreeLibraryAndExitThread),
	MHFF(kernel32, CreateFileA, HookFlags::HookCreateFile),
	MHFF(kernel32, CreateFileW, HookFlags::HookCreateFile),
	MHFF(kernel32, CreateThread, HookFlags::HookCreateThread),
};

static const int functionsCount = sizeof(functions) / sizeof(functions[0]);

HookedFunction* GetHookedFunction(void* originalFunction)
{
	static bool sorted = false;
	if(!sorted)
	{
		sort(functions, functions + functionsCount);
		sorted = true;
	}
	HookedFunction function;
	function.originalFunction = originalFunction;
	int functionId = lower_bound(functions, functions + functionsCount, function) - functions;
	if(functionId >= functionsCount || functions[functionId].originalFunction != function.originalFunction) return 0;

	if((callbackInformation->injectSettings.hookFlags & functions[functionId].hookFlags) != functions[functionId].hookFlags) return 0;

	return functions + functionId;
}
